import numpy as np
import platgo as pg


class ZDT5(pg.Problem):

    def __init__(self, D: int = 80) -> None:
        # 二进制编码，共十一个决策变量，x1为30位01编码，x2-x11均为5位，共计80位
        self.name = "ZDT5"
        self.type['multi'], self.type['binary'], self.type['large'], self.type['expensive'] = [True] * 4
        self.M = 2
        self.D = int(np.ceil(max(D - 30, 1)/5) * 5 + 30)
        self.borders = []
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        u = np.zeros((pop.decs.shape[0], 1 + int((pop.decs.shape[1]-30)/5)))
        u[:, 0] = np.sum(pop.decs[:, 0:30], 1)
        for i in range(1, u.shape[1]):
            u[:, i] = np.sum(pop.decs[:, (i-1) * 5 + 30:(i-1) * 5 + 35], axis=1)
        v = np.zeros(u.shape)
        v[u < 5] = 2 + u[u < 5]
        v[u == 5] = 1
        objv1 = 1 + u[:, 0]
        g = np.sum(v[:, 1:], 1)
        h = 1/objv1
        objv2 = g * h
        pop.objv = np.vstack((objv1, objv2)).T

    def get_optimal(self) -> np.ndarray:
        optimal1 = np.arange(1, 32)
        optimal2 = (self.D-30)/5/optimal1
        optimal = np.vstack((optimal1, optimal2)).T
        return optimal
